######################################################################
##
## Copyright (C) 2006,  Blekinge Institute of Technology
##
## Filename:      javamethod.py
## Author:        Simon Kagstrom <ska@bth.se>
## Description:   Java method stuff
##
## $Id: javamethod.py 14162 2007-03-11 14:00:55Z ska $
##
######################################################################
from Cibyl.BinaryTranslation import bytecode, register
from Cibyl.BinaryTranslation.Optimization.Mips import memoryRegisters
from Cibyl.BinaryTranslation.Optimization.Mips import mult
from Cibyl import config
import mips, instruction
from sets import Set

from Cibyl.BinaryTranslation.codeblock import CodeBlock

class JavaMethod(CodeBlock):
    """
    One method per function Java mehtod.
    """
    def __init__(self, function):
	self.function = function
	self.controller = function.controller
	self.labels = function.labels
	self.instructions = function.instructions
	self.bc = bytecode.ByteCodeGenerator(self.controller)
	self.rh = register.RegisterHandler(self.controller, self.bc)

	if function.address == self.controller.elf.getEntryPoint():
	    self.name = "start"
	    self.methodAccess = "public"
	else:
            sym = self.controller.elf.getSymbolByAddress(function.address)
            if sym.localLinkage():
                self.name = function.name +  "_%x" % (function.address)
            else:
                self.name = function.name
	    self.methodAccess = "private"

	CodeBlock.__init__(self, self.controller, function.instructions, function.labels, self.function.useTracing)

	# Each label and instruction belongs to one java method
	for lab in self.labels.values():
	    if lab.address >= self.address and lab.address <= self.address + self.size:
		lab.setJavaMethod(self)

	self.cleanupRegs = Set([mips.R_V0])
	if config.doRegisterScheduling:
	    self.cleanupRegs = Set()

	# For non-leafs, we always pass all registers
	if not self.function.isLeaf() or function.address == self.controller.elf.getEntryPoint():
	    self.addSourceRegister(mips.R_SP)
	    self.addSourceRegister(mips.R_A0)
	    self.addSourceRegister(mips.R_A1)
	    self.addSourceRegister(mips.R_A2)
	    self.addSourceRegister(mips.R_A3)
	    self.addDestinationRegister(mips.R_V0)
	    self.addDestinationRegister(mips.R_V1)

	for insn in self.instructions:
	    insn.setJavaMethod(self)

	self.argumentRegisters = [mips.R_SP, mips.R_A0, mips.R_A1, mips.R_A2, mips.R_A3]
	if not config.debug:
	    self.argumentRegisters = list(self.usedRegisters.intersection( [mips.R_SP, mips.R_A0, mips.R_A1, mips.R_A2, mips.R_A3] ))
	self.argumentRegisters.sort()
	# Always place sp first
	if mips.R_SP in self.argumentRegisters:
	    self.argumentRegisters.remove(mips.R_SP)
	    self.argumentRegisters = [mips.R_SP] + self.argumentRegisters

	# Perform skip stack store optimization. We actually always do
	# this since we that way avoid zeroing registers unecessary
	if not config.debug:
	    self.doSkipStackStoreOptimization()

	if config.doMemoryRegisterOptimization:
	    for bb in self.function.basicBlocks:
		memoryRegisters.run(bb)
		self.addDestinationRegisterSet(bb.destinationRegisters)
		self.addSourceRegisterSet(bb.sourceRegisters)

	if config.doMultOptimization:
	    mult.run(self.function)

    def getJavaMethodName(self):
	return self.name + "(" + 'I'*len(self.getRegistersToPass()) + ")" + self.getJavaReturnType()

    def getJavaReturnType(self):
	if not config.doRegisterScheduling:
	    return "I"
	if mips.R_V0 in self.usedRegisters:
	    return "I"
	else:
	    return "V"

    def getRegistersToPass(self):
	"Get the registers to pass to this function"
	if not config.doRegisterScheduling:
	    return [mips.R_SP, mips.R_A0, mips.R_A1, mips.R_A2, mips.R_A3]
	return self.argumentRegisters

    def doSkipStackStoreOptimization(self):
	"""Optimization to skip the stack saving and restoring since
	these are not actually needed with Cibyl (registers in local
	java variables)"""

	def instructionIsStackStore(insn):
	    return isinstance(insn, instruction.Sw) and insn.rt in mips.callerSavedRegisters and insn.rs == mips.R_SP

	def instructionIsStackLoad(insn):
	    return isinstance(insn, instruction.Lw) and insn.rt in mips.callerSavedRegisters and insn.rs == mips.R_SP

	for insn in self.function.basicBlocks[0].instructions:
	    if instructionIsStackStore(insn):
		if config.verbose: print "Removing", insn
		insn.nullify()

	returnBasicBlocks = self.function.getReturnBasicBlocks()
	for bb in returnBasicBlocks:
	    for insn in bb.instructions:
		# Validate that there is no use of the saved registers
		for dst in self.destinationRegisters:
		    if dst in mips.callerSavedRegisters and not instructionIsStackLoad(insn):
			return

	    for insn in bb.instructions:
		if instructionIsStackLoad(insn):
		    if config.verbose: print "Removing", insn
		    insn.nullify()

    def clobbersReg(self, reg):
	return reg in self.destinationRegisters

    def getNumberOfLocals(self):
	"Return the number of used registers by this function"
	if not config.doRegisterScheduling:
	    return mips.N_REGS
	return len(self.usedRegisters)


    def generateLookupTable(self):
	out = []
	for label in self.labels.values():
	    # Don't include the actual function address
	    if label.inJumpTab and label.address > self.address and label.address < self.address + self.size:
		cur = (label.address, str(label))
		out.append(cur)
	return out

    def fixup(self):
	"""This method constructs a list of clobbered registers
	(depending on the functions this function calls) and also adds
	registers to zero in the function."""
	# Add v1 to the list of clobbered registers if the destination
	# method clobbers V1
	for insn in self.instructions:
	    if isinstance(insn, instruction.Jal):
		otherMethod = self.controller.lookupJavaMethod(insn.dstAddress)
		if otherMethod.clobbersReg(mips.R_V1):
		    self.addDestinationRegister(mips.R_V1)
		if otherMethod.clobbersReg(mips.R_V0):
		    self.addDestinationRegister(mips.R_V0)
	    elif isinstance(insn, instruction.Jalr):
		self.addDestinationRegister(mips.R_V1)
		self.addDestinationRegister(mips.R_V0)

	# Somewhat conservative estimate of registers to zero - all
	# used registers which are not used as destinations in the
	# first basic block. This _requires_ that the s-register stack
	# removal optimization is turned on.
	for r in self.usedRegisters:
	    bb0 = self.function.basicBlocks[0]
	    # If debug is enabled, clear all local registers
	    if config.debug:
		if not register.staticRegs.has_key(r) and r not in self.argumentRegisters + [mips.R_ZERO]:
		    self.cleanupRegs.add(r)
		continue
	    if (r not in self.argumentRegisters + [mips.R_HI, mips.R_LO, mips.R_ZERO, mips.R_RA]) and (r not in bb0.destinationRegisters):
		self.cleanupRegs.add(r)

    def compile(self):
	"""Compile this method. The calling convention looks as
	follows:

	int method(int sp, int a0, int a1, int a2, int a3);

	The return result is register v0. Register v1 is placed in a
	static class-variable and is assigned after calls.
	"""
	self.controller.emit(".method %s static %s" % (self.methodAccess, self.getJavaMethodName()) )
	self.controller.emit(".limit stack 100")
	self.controller.emit(".limit locals %d" % self.getNumberOfLocals())

	# Generate a register mapping for this method and the registers to pass to the method
	if config.doRegisterScheduling:
	    registerMapping = register.generateRegisterMapping(self.usedRegisters, self.argumentRegisters, self.registerUseCount)

	    register.reg2local = registerMapping
	    for reg, local in registerMapping.iteritems():
		self.controller.emit(".var %2d is %s I from METHOD_START to METHOD_END" % (local, mips.registerNames[reg]))
	else:
	    for reg, local in register.reg2local.iteritems():
		if (reg not in self.usedRegisters and reg not in self.argumentRegisters) or reg in register.staticRegs:
		    continue
		self.controller.emit(".var %2d is %s I from METHOD_START to METHOD_END" % (local, mips.registerNames[reg]))

	for reg in self.cleanupRegs:
	    self.bc.pushConst(0)
	    self.rh.popToRegister(reg)

	if config.traceFunctionCalls:
	    self.bc.ldc("0x%08x: %s" % (self.address, self.name))
	    self.bc.invokestatic("CRunTime/emitFunctionEnterTrace(Ljava/lang/String;)V")

	self.controller.emit("METHOD_START:")
	# Compile this function
	self.function.compile()

	# Emit the local jumptable
	lookuptab = self.generateLookupTable()
	if lookuptab:
	    self.controller.emit("__CIBYL_local_jumptab:")
	    self.bc.lookupswitch(lookuptab, "__CIBYL_function_return")

	# Emit the return statement
	self.controller.emit("__CIBYL_function_return:")
	if config.traceFunctionCalls:
	    self.bc.ldc("")
	    self.bc.invokestatic("CRunTime/emitFunctionExitTrace(Ljava/lang/String;)V")
	if self.clobbersReg(mips.R_V1):
	    self.rh.pushRegister(mips.R_V1)
	    self.bc.putstatic("CRunTime/saved_v1 I")

	if self.getJavaReturnType() == "I":
	    self.rh.pushRegister(mips.R_V0)
	    self.bc.ireturn()
	else:
	    self.bc.emit_return()

	self.controller.emit("METHOD_END:")
	self.controller.emit(".end method")


class GlobalJumptabMethod(CodeBlock):
    def __init__(self, controller, methods):
	self.controller = controller
	self.name = "__CIBYL_global_jumptab(IIIIII)I"

	# Conservative estimate: produce a jumptab for all methods
	self.methods = methods

	self.bc = bytecode.ByteCodeGenerator(self.controller)
	self.rh = register.RegisterHandler(self.controller, self.bc)

	CodeBlock.__init__(self, self.controller, [], [], False)

    def generateLookupTable(self):
	out = []
	for method in self.methods:
	    cur = (method.address, "lab_" + method.name)
	    out.append(cur)
	return out

    def fixup(self):
	pass

    def compile(self):
	self.controller.emit(".method public static %s" % (self.name))
	self.controller.emit(".limit stack 100")
	self.controller.emit(".limit locals 8")

	# This is the register allocation for the global jumptab
	reg2local = {
	    mips.namesToRegisters["sp"] : 1,
	    mips.namesToRegisters["a0"] : 2,
	    mips.namesToRegisters["a1"] : 3,
	    mips.namesToRegisters["a2"] : 4,
	    mips.namesToRegisters["a3"] : 5,
	}

	if config.traceFunctionCalls:
	    self.bc.ldc("0x%08x: %s" % (self.address, self.name))
	    self.bc.invokestatic("CRunTime/emitFunctionEnterTrace(Ljava/lang/String;)V")

	self.bc.iload(0) # First argument is the method to call
	self.bc.pushConst(0)
	self.bc.istore(7)

	lookuptab = self.generateLookupTable()
	self.bc.lookupswitch(lookuptab, "error")

	for method in self.methods:
	    self.controller.emit("lab_" + method.name + ":")
	    for r in method.getRegistersToPass():
		self.bc.iload( reg2local[r])
	    self.bc.invokestatic("CompiledProgram/" + method.getJavaMethodName())
	    if method.getJavaReturnType() == "I":
		self.bc.istore(7)
	    self.bc.goto("out")

 	self.controller.emit("error:")
	self.bc.ldc("CIBYL ERROR: cannot find the address ")
	self.bc.iload(0)
	self.bc.invokestatic("CRunTime/emitTraceAddress(Ljava/lang/String;I)V")

 	self.controller.emit("out:")
	if config.traceFunctionCalls:
	    self.bc.ldc("")
	    self.bc.invokestatic("CRunTime/emitFunctionExitTrace(Ljava/lang/String;)V")
	self.bc.iload(7)
	self.bc.ireturn()
 	self.controller.emit(".end method")
